import os
import random


def train_valid_split(data_list):
    val_percent = 0.2
    random.shuffle(data_list)
    train_images_count = int((1 - val_percent) * len(data_list))
    valid_images_count = len(data_list) - train_images_count
    train_set = []
    for train_num in range(train_images_count):
        train_set.append(data_list[train_num])
    valid_set = []
    for valid_num in range(valid_images_count):
        valid_set.append(data_list[train_images_count + valid_num])
    return train_set, valid_set


def generate_label(datasets_dir):
    class_list = []
    image_list = []
    for _, direc, files in os.walk(datasets_dir):
        if len(direc) != 0:
            for ele in direc:
                class_list.append(ele)
        if len(files) != 0:
            image_list.append(files)
    all_obj = {}
    train_set = []
    valid_set = []
    for num in range(len(class_list)):
        all_obj[class_list[num]] = image_list[num]

    with open(datasets_dir + '/labels.txt', 'w') as f:
        for class_name in class_list:
            f.write(class_name + '\n')
    class_num = 0
    for class_name in all_obj.keys():
        train_data, valid_data = train_valid_split(all_obj[class_name])
        for train_obj in train_data:
            train_obj = class_name + '/' + train_obj + ' ' + str(class_num) + '\n'
            train_set.append(train_obj)
        for valid_obj in valid_data:
            valid_obj = class_name + '/' + valid_obj + ' ' + str(class_num) + '\n'
            valid_set.append(valid_obj)
        class_num += 1
    print(train_set)
    with open(datasets_dir + '/train_list.txt', 'w') as f:
        for train_label in train_set:
            f.write(train_label)
    with open(datasets_dir + '/val_list.txt', 'w') as f:
        for valid_label in valid_set:
            f.write(valid_label)
        # if file.endswith('.jpg') or file.endswith('.png'):
        #     label = root.split('/')[-1]
        #     note = label + '/' + file
        #     note = note + '\t' + label + '\n'
        #     f.write(note)
    print('[INFO]: Label gen.')


# def remove_path():
#     datasets_dir = './datasets/machine/gallery/'
#     with open(datasets_dir + 'label.txt', 'r') as f:
#         notes = f.readlines()
#
#     for note in notes:
#         path = note.split('\t')[0]
#         path = datasets_dir + path
#         if not os.path.exists(path):
#             notes.remove(note)
#
#     with open(datasets_dir + 'label.txt', 'w') as f:
#         for note in notes:
#             f.write(note)


if __name__ == "__main__":
    root_dir = '/home/stark/algo-env/datasets/smoking-calling'
    generate_label(root_dir)
    # for root, direc, files in os.walk(root_dir):
    #     print(root)
    #     print(direc)
    #     print(files)
    # gen_label(datasets_dir)
